/* This is sample from the OpenCV book. The copyright notice is below */

/* *************** License:**************************
   Oct. 3, 2008
   Right to use this code in any way you want without warranty, support or any
   guarantee of it working.

   BOOK: It would be nice if you cited it:
   Learning OpenCV: Computer Vision with the OpenCV Library
     by Gary Bradski and Adrian Kaehler
     Published by O'Reilly Media, October 3, 2008

   AVAILABLE AT:
     http://www.amazon.com/Learning-OpenCV-Computer-Vision-Library/dp/0596516134
     Or: http://oreilly.com/catalog/9780596516130/
     ISBN-10: 0596516134 or: ISBN-13: 978-0596516130

   OPENCV WEBSITES:
     Homepage:      http://opencv.org
     Online docs:   http://docs.opencv.org
     Q&A forum:     http://answers.opencv.org
     Issue tracker: http://code.opencv.org
     GitHub:        https://github.com/opencv/opencv/
   ************************************************** */

#include "opencv2/calib3d/calib3d.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"

#include <algorithm>
#include <ctype.h>
#include <iostream>
#include <iterator>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <vector>

using namespace cv;
using namespace std;

static int print_help() {
  cout << "Usage:\n ./stereo_proc  <image list XML/YML file>\n" << endl;
  return 0;
}

static void StereoCalib(const vector<string> &imagelist, Size boardSize,
                        bool useCalibrated = true, bool showRectified = true) {
  if (imagelist.size() % 2 != 0) {
    cout
        << "Error: the image list contains odd (non-even) number of elements\n";
    return;
  }

  bool displayCorners = false; // true;
  const int maxScale = 2;
  const float squareSize = 1.f; // Set this to your actual square size
  // ARRAY AND VECTOR STORAGE:

  vector<vector<Point2f> > imagePoints[2];
  vector<vector<Point3f> > objectPoints;
  Size imageSize;

  int i, j, k, nimages = (int)imagelist.size() / 2;

  imagePoints[0].resize(nimages);
  imagePoints[1].resize(nimages);
  vector<string> goodImageList;

  for (i = j = 0; i < nimages; i++) {
    for (k = 0; k < 2; k++) {
      const string &filename = imagelist[i * 2 + k];
      Mat img = imread(filename, 0);
      if (img.empty())
        break;
      if (imageSize == Size())
        imageSize = img.size();
      else if (img.size() != imageSize) {
        cout << "The image " << filename
             << " has the size different from the first image size. Skipping "
                "the pair\n";
        break;
      }
    }
    if (k == 2) {
      goodImageList.push_back(imagelist[i * 2]);
      goodImageList.push_back(imagelist[i * 2 + 1]);
      j++;
    }
  }
  cout << j << " pairs have been successfully detected.\n";
  nimages = j;
  if (nimages < 2) {
    cout << "Error: too little pairs to run the calibration\n";
    return;
  }

  Mat cameraMatrix[2], distCoeffs[2];
  cameraMatrix[0] = Mat::eye(3, 3, CV_64F);
  cameraMatrix[1] = Mat::eye(3, 3, CV_64F);
  Mat R, T, E, F;

  // save intrinsic parameters
  FileStorage fs("intrinsics.yml", CV_STORAGE_READ);
  FileNode n = fs["strings"];
  if (fs.isOpened()) {

    fs["M1"] >> cameraMatrix[0];
    fs["D1"] >> distCoeffs[0];
    fs["M2"] >> cameraMatrix[1];
    fs["D2"] >> distCoeffs[1];
    fs.release();
  } else
    cout << "Error: can not save the intrinsic parameters\n";

  Mat R1, R2, P1, P2, Q;
  Rect validRoi[2];

  fs.open("extrinsics.yml", CV_STORAGE_READ);
  if (fs.isOpened()) {
    cv::String st;
    fs["R"] >> R;
    fs["T"] >> T;
    fs["R1"] >> R1;
    fs["R2"] >> R2;
    fs["P1"] >> P1;
    fs["P2"] >> P2;
    fs["Q"] >> Q;
    fs.release();
  } else
    cout << "Error: can not save the extrinsic parameters\n";

  stereoRectify(cameraMatrix[0], distCoeffs[0], cameraMatrix[1], distCoeffs[1],
                imageSize, R, T, R1, R2, P1, P2, Q, CALIB_ZERO_DISPARITY, 1,
                imageSize, &validRoi[0], &validRoi[1]);

  // OpenCV can handle left-right
  // or up-down camera arrangements
  bool isVerticalStereo = fabs(P2.at<double>(1, 3)) > fabs(P2.at<double>(0, 3));

  // COMPUTE AND DISPLAY RECTIFICATION
  if (!showRectified)
    return;

  Mat rmap[2][2];
  // IF BY CALIBRATED (BOUGUET'S METHOD)
  if (useCalibrated) {
    // we already computed everything
  }
  // OR ELSE HARTLEY'S METHOD
  else
  // use intrinsic parameters of each camera, but
  // compute the rectification transformation directly
  // from the fundamental matrix
  {
    vector<Point2f> allimgpt[2];
    for (k = 0; k < 2; k++) {
      for (i = 0; i < nimages; i++)
        std::copy(imagePoints[k][i].begin(), imagePoints[k][i].end(),
                  back_inserter(allimgpt[k]));
    }
    F = findFundamentalMat(Mat(allimgpt[0]), Mat(allimgpt[1]), FM_8POINT, 0, 0);
    Mat H1, H2;
    stereoRectifyUncalibrated(Mat(allimgpt[0]), Mat(allimgpt[1]), F, imageSize,
                              H1, H2, 3);

    R1 = cameraMatrix[0].inv() * H1 * cameraMatrix[0];
    R2 = cameraMatrix[1].inv() * H2 * cameraMatrix[1];
    P1 = cameraMatrix[0];
    P2 = cameraMatrix[1];
  }

  // Precompute maps for cv::remap()
  initUndistortRectifyMap(cameraMatrix[0], distCoeffs[0], R1, P1, imageSize,
                          CV_16SC2, rmap[0][0], rmap[0][1]);
  initUndistortRectifyMap(cameraMatrix[1], distCoeffs[1], R2, P2, imageSize,
                          CV_16SC2, rmap[1][0], rmap[1][1]);

  Mat canvas;
  double sf;
  int w, h;
  if (!isVerticalStereo) {
    sf = 600. / MAX(imageSize.width, imageSize.height);
    w = cvRound(imageSize.width * sf);
    h = cvRound(imageSize.height * sf);
    canvas.create(h, w * 2, CV_8UC3);
  } else {
    sf = 300. / MAX(imageSize.width, imageSize.height);
    w = cvRound(imageSize.width * sf);
    h = cvRound(imageSize.height * sf);
    canvas.create(h * 2, w, CV_8UC3);
  }

  for (i = 0; i < nimages; i++) {
    for (k = 0; k < 2; k++) {
      Mat img = imread(goodImageList[i * 2 + k], 0), rimg, cimg;
      remap(img, rimg, rmap[k][0], rmap[k][1], CV_INTER_LINEAR);
      cvtColor(rimg, cimg, COLOR_GRAY2BGR);
      Mat canvasPart = !isVerticalStereo ? canvas(Rect(w * k, 0, w, h))
                                         : canvas(Rect(0, h * k, w, h));
      resize(cimg, canvasPart, canvasPart.size(), 0, 0, CV_INTER_AREA);
      if (useCalibrated) {
        Rect vroi(cvRound(validRoi[k].x * sf), cvRound(validRoi[k].y * sf),
                  cvRound(validRoi[k].width * sf),
                  cvRound(validRoi[k].height * sf));
        rectangle(canvasPart, vroi, Scalar(0, 0, 255), 3, 8);
      }
    }

    if (!isVerticalStereo)
      for (j = 0; j < canvas.rows; j += 16)
        line(canvas, Point(0, j), Point(canvas.cols, j), Scalar(0, 255, 0), 1,
             8);
    else
      for (j = 0; j < canvas.cols; j += 16)
        line(canvas, Point(j, 0), Point(j, canvas.rows), Scalar(0, 255, 0), 1,
             8);
    imshow("rectified", canvas);
    char c = (char)waitKey();
    if (c == 27 || c == 'q' || c == 'Q')
      break;
  }
}

static bool readStringList(const string &filename, vector<string> &l) {
  l.resize(0);
  FileStorage fs(filename, FileStorage::READ);
  if (!fs.isOpened())
    return false;
  FileNode n = fs.getFirstTopLevelNode();
  if (n.type() != FileNode::SEQ)
    return false;
  FileNodeIterator it = n.begin(), it_end = n.end();
  for (; it != it_end; ++it)
    l.push_back((string)*it);
  return true;
}

int main(int argc, char **argv) {
  Size boardSize;
  string imagelistfn;
  bool showRectified = true;

  imagelistfn = argv[1];

  if (imagelistfn == "") {
    imagelistfn = "stereo_calib.xml";
    boardSize = Size(9, 6);
  }

  vector<string> imagelist;
  bool ok = readStringList(imagelistfn, imagelist);
  if (!ok || imagelist.empty()) {
    cout << "can not open " << imagelistfn << " or the string list is empty"
         << endl;
    return print_help();
  }

  StereoCalib(imagelist, boardSize, true, showRectified);
  return 0;
}
